
while(T) {

  cat(
    "---------------------------------------------------------------------\n",
    "Iteration ", iter, "\n\n", sep=""
  )
  # update q
  q_new <- update_q(
    theta = theta, # to calculate the log likelihood
    v = vx, # hid, s are columns of vx
    full_data = full_data,
    state.vars,
    pi,
    q,
    base_path = base_path
  )

  # q0, Q, QY
  q0 <- q_new[! duplicated(q_new[,1]),c(1,3:ncol(q_new))]
  Q <- q_new[,3:ncol(q_new)]
  for(s in 1:ncol(Q)) QY[,s] <- rrf$Y.orig * Q[,s]
  QY <- Matrix(QY, sparse=T)
  Q <- as.matrix(Q)

  # update pi
  pi_new <- update_pi(q_new)

  # # # # # #
  # update psi (rest of v does not change)
  # # # # # #
  old_psi <- vx[,"psi"]

  # This starts the update of psi if the psi_parallel%d.R workers are running
  #  in the background.
  for(s in 1:S) {
    saveRDS(q_new[,sprintf("state%02d",s)], sprintf("%s/iter/q%d.RDS", base_path, s))
    saveRDS(0, sprintf("%s/iter/go_files/go_file%d", base_path, s))
  }
  # check if all procs finished: does finish flag exist for all s?
  all_flags <- F
  while(!all_flags) {
    Sys.sleep(60)
    all_flags <- T
    for(s in 1:S) all_flags <- all_flags & file.exists(sprintf("%s/iter/flag%d", base_path, s))
  }
  for(s in 1:S) {
    file.remove(sprintf("%s/iter/flag%d", base_path, s))
    # # # #
    # update psi in vx[,"psi"] for each s
    new_psi <- readRDS(sprintf("%s/iter/psi%d.RDS", base_path, s)) # order: hid > year > choice / check this!
    vx[vx[,"s"] == s,"psi"] <- new_psi[,4]
    #
    file.remove(sprintf("%s/iter/psi%d.RDS", base_path, s))
    # # # #
  }
  # # # # # #

  # store psi of current iteration (in case the algorithm crashes)
  saveRDS(vx[,c("hid", "year", "depvar", "psi", "s")], sprintf("%s/iter/psi_iter%03d.RDS", base_path, iter))

  # update theta
  fit_new <- update_theta(
    v = vx,
    data = data,
    theta = theta,
    q0 = q0,
    J = max(v[,"depvar"]),
    start_index = start_index
  )
  theta_new <- fit_new$theta
  LL_new <- fit_new$LogLik

  #
  if(length(pi) > 1) {
    d.theta <- sqrt(mean((theta_new - theta)^2))
    d.pi <- sqrt(mean((pi_new - pi)^2))
    d.LL <- LL_new - LL
    cat(
      "\n",
      "pi = ", sprintf("%.4f ", pi_new), "\n",
      "d.theta = ", sprintf("%.6f ", d.theta), ", ",
      "d.pi = ", sprintf("%.6f", d.pi), "\n",
      sprintf("LL = %.2f, d.LL = %.6f", LL_new, d.LL), "\n",
      "---------------------------------------------------------------------\n\n",
      sep=""
    )
    print(theta_new)

    if(d.LL < 1e-6) {
      cat(
        "\n\n---------------------------------------------------------------------\n\n",
        "\nEM Algorithm converged to local maximum.\n",
        "---------------------------------------------------------------------\n",
        sep=""
      );
      saveRDS(
        list(
          res = fit,
          q = q_new,
          pi = pi_new,
          Q = Q,
          QY = QY,
          v = vx,
          psi = old_psi,
          iter = iter
        ),
        sprintf("%s/iter/iter_tmp_res.RDS", base_path)
      )
      file.remove(sprintf("%s/iter/go_files/keep_going", base_path))
      break;
    } else {
      saveRDS(
        list(
          res = fit_new,
          q = q_new,
          pi = pi_new,
          psi = vx[,"psi"],
          old_psi = old_psi,
          iter = iter
        ),
        sprintf("%s/iter/iter_%03d_res.RDS", base_path, iter)
      )
    }

    # update values
    LL <- LL_new
    theta <- theta_new
    q <- q_new
    pi <- pi_new
    fit <- fit_new

    # increment iter counter
    iter <- iter + 1
  } else {
    break;
  }

}
